{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Overview\n",
    "\n",
    "This notebook calculates and plots: \n",
    "* $E_z$ field distribution at time_start and time_end;\n",
    "* total EM energy evolution in time;\n",
    "* NCI growth rate, $Im( \\omega )/\\omega_{p,r}$, calculated between time_start and time_end (within the linear growth of the total energy).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import scipy.constants as scc\n",
    "import yt\n",
    "from scipy.constants import c\n",
    "\n",
    "yt.funcs.mylog.setLevel(50)\n",
    "import glob\n",
    "\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Plot NCI growth rate\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "path_wx = \"path to diags folder\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "file_list_warpx = glob.glob(path_wx + \"diag1?????\")\n",
    "iterations_warpx = [\n",
    "    int(file_name[len(file_name) - 5 :]) for file_name in file_list_warpx\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def calculate_parameters(path):\n",
    "    iteration = 200\n",
    "    dsx = yt.load(path + \"diag1%05d/\" % iteration)\n",
    "    dxx = dsx.domain_width / dsx.domain_dimensions\n",
    "    dx = dxx[0]\n",
    "    dx = 1.0 * dx.ndarray_view()\n",
    "\n",
    "    dz = dxx[1]\n",
    "\n",
    "    dz = 1.0 * dz.ndarray_view()\n",
    "\n",
    "    ds1 = yt.load(path + \"/diag100100/\")\n",
    "    ds2 = yt.load(path + \"/diag100200/\")\n",
    "\n",
    "    cur_t1 = ds1.current_time\n",
    "    cur_t2 = ds2.current_time\n",
    "    cur_t2.to_ndarray\n",
    "    dt = (cur_t2 - cur_t1) / 100\n",
    "    dt = 1.0 * dt.ndarray_view()\n",
    "    return dx, dz, dt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dx, dz, dt = calculate_parameters(path_wx)\n",
    "print(dx, dz, dt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def get_fourier_transform_wx(path, fieldcomp, iteration, plot=False, remove_last=True):\n",
    "    \"\"\"\n",
    "    Calculate the Fourier transform of the field at a given iteration\n",
    "    \"\"\"\n",
    "\n",
    "    ds = yt.load(path + \"/diag1%05d/\" % iteration)\n",
    "\n",
    "    grid = ds.index.grids[0]\n",
    "    F = grid[fieldcomp]\n",
    "    F = F.ndarray_view()\n",
    "\n",
    "    if remove_last:\n",
    "        F = F[:-1, :-1]\n",
    "    F = F[:, :, 0]\n",
    "\n",
    "    kxmax = np.pi / dx\n",
    "    kzmax = np.pi / dz\n",
    "    Nx = F.shape[0]\n",
    "    Nz = F.shape[1]\n",
    "    spectralF = np.fft.fftshift(np.fft.fft2(F))[int(Nx / 2) :, int(Nz / 2) :]\n",
    "\n",
    "    if plot:\n",
    "        plt.imshow(np.log(abs(spectralF)), origin=\"lower\", extent=[0, kxmax, 0, kzmax])\n",
    "        plt.colorbar()\n",
    "\n",
    "    return (spectralF, kxmax, kzmax)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def growth_rate_between_wx(\n",
    "    path, iteration1, iteration2, remove_last=False, threshold=-13\n",
    "):\n",
    "    \"\"\"\n",
    "    Calculate the difference in spectral amplitude between two iterations,\n",
    "    return the growth rate\n",
    "\n",
    "    \"\"\"\n",
    "    spec1, kxmax, kzmax = get_fourier_transform_wx(\n",
    "        path, \"Ez\", iteration=iteration1, remove_last=remove_last\n",
    "    )\n",
    "    spec1 = np.where(abs(spec1) > np.exp(threshold), spec1, np.exp(threshold))\n",
    "\n",
    "    spec2, kxmax, kzmax = get_fourier_transform_wx(\n",
    "        path, \"Ez\", iteration=iteration2, remove_last=remove_last\n",
    "    )\n",
    "\n",
    "    spec2 = np.where(abs(spec2) > np.exp(threshold), spec2, np.exp(threshold))\n",
    "    diff_growth = np.log(abs(spec2)) - np.log(abs(spec1))\n",
    "\n",
    "    diff_time = (iteration2 - iteration1) * dt\n",
    "    growth_rate = diff_growth / diff_time / c\n",
    "\n",
    "    return (growth_rate, [0, kxmax, 0, kzmax])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def energy(ts):\n",
    "    Ex = ts.index.grids[0][\"boxlib\", \"Ex\"].squeeze().v\n",
    "    Ey = ts.index.grids[0][\"boxlib\", \"Ey\"].squeeze().v\n",
    "    Ez = ts.index.grids[0][\"boxlib\", \"Ez\"].squeeze().v\n",
    "\n",
    "    Bx = ts.index.grids[0][\"boxlib\", \"Ex\"].squeeze().v\n",
    "    By = ts.index.grids[0][\"boxlib\", \"Ey\"].squeeze().v\n",
    "    Bz = ts.index.grids[0][\"boxlib\", \"Ez\"].squeeze().v\n",
    "\n",
    "    energyE = scc.epsilon_0 * np.sum(Ex**2 + Ey**2 + Ez**2)\n",
    "    energyB = np.sum(Bx**2 + By**2 + Bz**2) / scc.mu_0\n",
    "    energy = energyE + energyB\n",
    "    return energy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "energy_list = []\n",
    "for iter in iterations_warpx:\n",
    "    path = path_wx + \"/diag1%05d/\" % iter\n",
    "    ds = yt.load(path)\n",
    "    energy_list.append(energy(ds))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "iteration_start = 1700\n",
    "iteration_end = 1900\n",
    "iter_delta = iterations_warpx[2] - iterations_warpx[1]\n",
    "\n",
    "\n",
    "ds_start = yt.load(path_wx + \"diag1%05d/\" % iteration_start)\n",
    "Ez_start = ds.index.grids[0][\"boxlib\", \"Ez\"].squeeze().v\n",
    "\n",
    "ds_end = yt.load(path_wx + \"diag1%05d/\" % iteration_end)\n",
    "Ez_end = ds_end.index.grids[0][\"boxlib\", \"Ez\"].squeeze().v\n",
    "\n",
    "gr_wx, extent = growth_rate_between_wx(path_wx, iteration_start, iteration_end)\n",
    "\n",
    "\n",
    "fs = 13\n",
    "vmax = 0.05\n",
    "vmin = -vmax\n",
    "cmap_special = \"bwr\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, (ax1, ax2) = plt.subplots(ncols=2, figsize=(10, 4.0))\n",
    "\n",
    "fs = 14\n",
    "cmap = \"viridis\"\n",
    "aspect = \"auto\"\n",
    "\n",
    "img1 = ax1.imshow(Ez_start, aspect=aspect, cmap=cmap, extent=extent)\n",
    "ax1.set_title(\"$t_{step}=$%i\" % iteration_start, size=fs)\n",
    "ax1.set_xlabel(\" $k_{p,r} z$ \", size=fs)\n",
    "ax1.set_ylabel(\" $k_{p,r} x $ \", size=fs)\n",
    "fig.colorbar(img1, ax=ax1, label=\"$E_z, [V/m]$\")\n",
    "\n",
    "img2 = ax2.imshow(Ez_end, aspect=aspect, cmap=cmap, extent=extent)\n",
    "ax2.set_title(\"$t_{step}=$%i\" % iteration_end, size=fs)\n",
    "ax2.set_xlabel(\" $k_{p,r} z$ \", size=fs)\n",
    "ax2.set_ylabel(\" $k_{p,r} x $ \", size=fs)\n",
    "fig.colorbar(img2, ax=ax2, label=\"$E_z, [V/m]$\")\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, (ax1, ax2) = plt.subplots(ncols=2, nrows=1, figsize=(13, 5.0))\n",
    "\n",
    "fs = 14\n",
    "\n",
    "img1 = ax1.semilogy(iterations_warpx, energy_list)\n",
    "ax1.semilogy(iteration_start, energy_list[iteration_start / iter_delta], \"ro\")\n",
    "ax1.semilogy(iteration_end, energy_list[iteration_end / iter_delta], \"ro\")\n",
    "ax1.grid()\n",
    "ax1.legend()\n",
    "ax1.set_xlabel(\"time step\", size=fs)\n",
    "ax1.set_ylabel(\"Total EM energy\", size=fs)\n",
    "\n",
    "img2 = ax2.imshow(\n",
    "    gr_wx,\n",
    "    origin=\"lower\",\n",
    "    cmap=\"bwr\",\n",
    "    vmax=0.05,\n",
    "    vmin=-vmax,\n",
    "    interpolation=\"nearest\",\n",
    "    extent=[0, 1, 0, 1],\n",
    ")\n",
    "ax2.set_title(\"NCI growth rate\", size=fs)\n",
    "ax2.set_xlabel(\"$k_{p,r} z$ \", size=fs)\n",
    "ax2.set_ylabel(\"$k_{p,r} x $ \", size=fs)\n",
    "fig.colorbar(img2, ax=ax2, label=\"$Im(\\omega)/\\omega_{p,r}$\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
